package dns

import (
	"context"
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/catroll/fastdns/internal/config"
	"github.com/catroll/fastdns/internal/logger"
)

// Resolver 是 DNS 解析器的主要结构体
type Resolver struct {
	config      *config.Config
	cache       *Cache
	upstreams   []*UpstreamServer
	upstreamsMu sync.RWMutex
}

// UpstreamServer 表示上游 DNS 服务器
type UpstreamServer struct {
	Address  string
	Protocol string
	Timeout  time.Duration
	Priority int
	client   *net.Resolver
	healthy  bool
	latency  time.Duration
}

func (us *UpstreamServer) Healthy() bool {
	return us.healthy
}

func (us *UpstreamServer) Latency() time.Duration {
	return us.latency
}

// NewResolver 创建一个新的 DNS 解析器
func NewResolver(cfg *config.Config) (*Resolver, error) {
	r := &Resolver{
		config:    cfg,
		upstreams: make([]*UpstreamServer, 0, len(cfg.DNS.Upstream)),
	}

	// 初始化缓存
	if cfg.DNS.Cache.Enabled {
		r.cache = NewCache(cfg.DNS.Cache.MaxSize, cfg.DNS.Cache.TTL)
		logger.Info("DNS 缓存已启用", "max_size", cfg.DNS.Cache.MaxSize, "ttl", cfg.DNS.Cache.TTL)
	} else {
		logger.Info("DNS 缓存已禁用")
	}

	// 初始化上游 DNS 服务器
	for _, us := range cfg.DNS.Upstream {
		upstream := &UpstreamServer{
			Address:  us.Address,
			Protocol: us.Protocol,
			Timeout:  us.Timeout,
			Priority: us.Priority,
			healthy:  true, // 初始假设所有服务器都是健康的
		}

		// 创建自定义解析器
		upstream.client = &net.Resolver{
			PreferGo: true,
			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
				d := net.Dialer{
					Timeout: upstream.Timeout,
				}
				return d.DialContext(ctx, upstream.Protocol, upstream.Address)
			},
		}

		r.upstreams = append(r.upstreams, upstream)
		logger.Info("添加上游 DNS 服务器",
			"address", upstream.Address,
			"protocol", upstream.Protocol,
			"timeout", upstream.Timeout,
			"priority", upstream.Priority)
	}

	// 按优先级排序上游服务器
	r.sortUpstreams()

	return r, nil
}

// sortUpstreams 按优先级排序上游服务器
func (r *Resolver) sortUpstreams() {
	r.upstreamsMu.Lock()
	defer r.upstreamsMu.Unlock()

	// 按优先级排序（优先级数字越小越优先）
	for i := 0; i < len(r.upstreams); i++ {
		for j := i + 1; j < len(r.upstreams); j++ {
			if r.upstreams[i].Priority > r.upstreams[j].Priority {
				r.upstreams[i], r.upstreams[j] = r.upstreams[j], r.upstreams[i]
			}
		}
	}
}

// LookupIP 查询域名对应的 IP 地址
func (r *Resolver) LookupIP(ctx context.Context, domain string) ([]net.IP, error) {
	// 确定记录类型，默认为 A 记录
	recordType := RTypeA

	logger.Info("开始 DNS 查询", "domain", domain, "type", recordType)
	startTime := time.Now()

	// 检查缓存
	if r.cache != nil {
		// 先尝试查询 A 记录缓存
		if ips, found := r.cache.GetIPs(domain, RTypeA); found {
			logger.Info("DNS A 记录缓存命中", "domain", domain, "ips", ips, "duration", time.Since(startTime))
			return ips, nil
		}

		// 如果未找到 A 记录，尝试查询 AAAA 记录
		if ips, found := r.cache.GetIPs(domain, RTypeAAAA); found {
			logger.Info("DNS AAAA 记录缓存命中", "domain", domain, "ips", ips, "duration", time.Since(startTime))
			return ips, nil
		}

		logger.Debug("DNS 缓存未命中", "domain", domain)
	}

	// 查询上游 DNS 服务器
	logger.Debug("从上游 DNS 服务器查询", "domain", domain)
	ips, err := r.lookupIPFromUpstreams(ctx, domain)
	if err != nil {
		logger.Error("DNS 查询失败", "domain", domain, "error", err, "duration", time.Since(startTime))
		return nil, err
	}

	// 更新缓存
	if r.cache != nil && len(ips) > 0 {
		// 分类存储 IPv4 和 IPv6 地址
		ipv4s := make([]net.IP, 0)
		ipv6s := make([]net.IP, 0)

		for _, ip := range ips {
			if ip.To4() != nil {
				ipv4s = append(ipv4s, ip)
			} else {
				ipv6s = append(ipv6s, ip)
			}
		}

		// 分别缓存 A 和 AAAA 记录
		if len(ipv4s) > 0 {
			r.cache.SetIPs(domain, RTypeA, ipv4s)
			logger.Debug("DNS A 记录缓存更新", "domain", domain, "ips", ipv4s)
		}

		if len(ipv6s) > 0 {
			r.cache.SetIPs(domain, RTypeAAAA, ipv6s)
			logger.Debug("DNS AAAA 记录缓存更新", "domain", domain, "ips", ipv6s)
		}
	}

	logger.Info("DNS 查询完成", "domain", domain, "ips", ips, "count", len(ips), "duration", time.Since(startTime))
	return ips, nil
}

// lookupIPFromUpstreams 从上游 DNS 服务器查询 IP 地址
func (r *Resolver) lookupIPFromUpstreams(ctx context.Context, domain string) ([]net.IP, error) {
	logger.Debug("从上游 DNS 服务器查询", "domain", domain)

	// 使用读锁以避免在查询过程中修改上游服务器列表
	r.upstreamsMu.RLock()
	defer r.upstreamsMu.RUnlock()

	// 检查查询模式
	if r.config.DNS.QueryMode == "all" {
		// 使用“所有”模式 - 查询所有健康的上游服务器并合并结果
		return r.lookupIPFromAllUpstreams(ctx, domain)
	} else {
		// 使用“随机”模式 - 使用第一个成功的结果
		return r.lookupIPFromRandomUpstream(ctx, domain)
	}
}

// lookupIPFromRandomUpstream 从上游 DNS 服务器随机获取一个成功的查询结果
func (r *Resolver) lookupIPFromRandomUpstream(ctx context.Context, domain string) ([]net.IP, error) {
	// 创建通道以接收结果和错误
	resultCh := make(chan []net.IP, len(r.upstreams))
	errCh := make(chan error, len(r.upstreams))

	// 并行查询所有上游服务器
	for _, upstream := range r.upstreams {
		if !upstream.healthy {
			logger.Debug("跳过不健康的上游 DNS 服务器", "address", upstream.Address)
			errCh <- fmt.Errorf("上游 DNS 服务器不健康: %s", upstream.Address)
			continue
		}

		go func(us *UpstreamServer) {
			// 创建带超时的上下文
			timeoutCtx, cancel := context.WithTimeout(ctx, us.Timeout)
			defer cancel()

			// 记录查询开始时间
			start := time.Now()

			// 查询 IP 地址
			ips, err := us.client.LookupIP(timeoutCtx, "ip", domain)
			latency := time.Since(start)

			// 更新延迟统计
			us.latency = latency

			if err != nil {
				logger.Debug("上游 DNS 查询失败",
					"address", us.Address,
					"domain", domain,
					"error", err,
					"latency", latency)
				errCh <- err
			} else {
				logger.Debug("上游 DNS 查询成功",
					"address", us.Address,
					"domain", domain,
					"ips_count", len(ips),
					"latency", latency)
				resultCh <- ips
			}
		}(upstream)
	}

	// 等待第一个成功的结果或所有错误
	for i := 0; i < len(r.upstreams); i++ {
		select {
		case ips := <-resultCh:
			return ips, nil
		case err := <-errCh:
			// 继续等待其他结果
			if i == len(r.upstreams)-1 {
				// 如果所有上游服务器都失败，返回最后一个错误
				return nil, fmt.Errorf("所有上游 DNS 服务器查询失败: %w", err)
			}
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}

	return nil, fmt.Errorf("所有上游 DNS 服务器查询失败")
}

// lookupIPFromAllUpstreams 从所有上游 DNS 服务器获取结果并合并
func (r *Resolver) lookupIPFromAllUpstreams(ctx context.Context, domain string) ([]net.IP, error) {
	// 创建通道以接收结果和错误
	type resultType struct {
		ips []net.IP
		err error
		us  *UpstreamServer
	}
	resultCh := make(chan resultType, len(r.upstreams))

	// 记录健康的上游服务器数量
	healthyCount := 0
	for _, upstream := range r.upstreams {
		if upstream.healthy {
			healthyCount++
		}
	}

	// 如果没有健康的上游服务器，直接返回错误
	if healthyCount == 0 {
		return nil, fmt.Errorf("没有健康的上游 DNS 服务器")
	}

	// 并行查询所有健康的上游服务器
	for _, upstream := range r.upstreams {
		if !upstream.healthy {
			logger.Debug("跳过不健康的上游 DNS 服务器", "address", upstream.Address)
			continue
		}

		go func(us *UpstreamServer) {
			// 创建带超时的上下文
			timeoutCtx, cancel := context.WithTimeout(ctx, us.Timeout)
			defer cancel()

			// 记录查询开始时间
			start := time.Now()

			// 查询 IP 地址
			ips, err := us.client.LookupIP(timeoutCtx, "ip", domain)
			latency := time.Since(start)

			// 更新延迟统计
			us.latency = latency

			if err != nil {
				logger.Debug("上游 DNS 查询失败",
					"address", us.Address,
					"domain", domain,
					"error", err,
					"latency", latency)
				resultCh <- resultType{nil, err, us}
				return
			}

			logger.Debug("上游 DNS 查询成功",
				"address", us.Address,
				"domain", domain,
				"ips_count", len(ips),
				"latency", latency)
			resultCh <- resultType{ips, nil, us}
		}(upstream)
	}

	// 收集所有结果
	allIPs := make(map[string]net.IP)
	successCount := 0
	failCount := 0
	var lastErr error

	// 等待所有查询完成
	for i := 0; i < healthyCount; i++ {
		select {
		case result := <-resultCh:
			if result.err != nil {
				failCount++
				lastErr = result.err
			} else {
				successCount++
				// 将 IP 地址添加到结果集合中，使用字符串表示作为键以去重
				for _, ip := range result.ips {
					allIPs[ip.String()] = ip
				}
			}
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}

	// 如果没有成功的查询，返回错误
	if successCount == 0 {
		return nil, fmt.Errorf("所有上游 DNS 服务器查询失败: %w", lastErr)
	}

	// 将所有唯一的 IP 地址转换为切片
	uniqIPs := make([]net.IP, 0, len(allIPs))
	for _, ip := range allIPs {
		uniqIPs = append(uniqIPs, ip)
	}

	logger.Debug("合并上游 DNS 查询结果",
		"domain", domain,
		"total_upstreams", healthyCount,
		"success_upstreams", successCount,
		"fail_upstreams", failCount,
		"unique_ips", len(uniqIPs))

	return uniqIPs, nil
}

// AddUpstreamServer 添加新的上游 DNS 服务器
func (r *Resolver) AddUpstreamServer(address, protocol string, timeout time.Duration, priority int) error {
	r.upstreamsMu.Lock()
	defer r.upstreamsMu.Unlock()

	// 检查是否已存在
	for _, us := range r.upstreams {
		if us.Address == address {
			return fmt.Errorf("上游 DNS 服务器已存在: %s", address)
		}
	}

	// 创建新的上游服务器
	upstream := &UpstreamServer{
		Address:  address,
		Protocol: protocol,
		Timeout:  timeout,
		Priority: priority,
		healthy:  true,
	}

	// 创建自定义解析器
	upstream.client = &net.Resolver{
		PreferGo: true,
		Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
			d := net.Dialer{
				Timeout: upstream.Timeout,
			}
			return d.DialContext(ctx, upstream.Protocol, upstream.Address)
		},
	}

	// 添加到列表
	r.upstreams = append(r.upstreams, upstream)

	// 重新排序
	r.sortUpstreams()

	logger.Info("添加上游 DNS 服务器",
		"address", upstream.Address,
		"protocol", upstream.Protocol,
		"timeout", upstream.Timeout,
		"priority", upstream.Priority)

	return nil
}

// RemoveUpstreamServer 移除上游 DNS 服务器
func (r *Resolver) RemoveUpstreamServer(address string) error {
	r.upstreamsMu.Lock()
	defer r.upstreamsMu.Unlock()

	// 至少保留一个上游服务器
	if len(r.upstreams) <= 1 {
		return fmt.Errorf("无法移除唯一的上游 DNS 服务器")
	}

	// 查找并移除
	for i, us := range r.upstreams {
		if us.Address == address {
			r.upstreams = append(r.upstreams[:i], r.upstreams[i+1:]...)
			logger.Info("移除上游 DNS 服务器", "address", address)
			return nil
		}
	}

	return fmt.Errorf("上游 DNS 服务器不存在: %s", address)
}

// GetUpstreamServers 获取所有上游 DNS 服务器
func (r *Resolver) GetUpstreamServers() []*UpstreamServer {
	r.upstreamsMu.RLock()
	defer r.upstreamsMu.RUnlock()

	// 返回副本以避免并发问题
	result := make([]*UpstreamServer, len(r.upstreams))
	copy(result, r.upstreams)
	return result
}

// CheckUpstreamHealth 检查上游 DNS 服务器的健康状态
func (r *Resolver) CheckUpstreamHealth() {
	r.upstreamsMu.RLock()
	upstreams := make([]*UpstreamServer, len(r.upstreams))
	copy(upstreams, r.upstreams)
	r.upstreamsMu.RUnlock()

	for _, us := range upstreams {
		go func(upstream *UpstreamServer) {
			ctx, cancel := context.WithTimeout(context.Background(), upstream.Timeout)
			defer cancel()

			start := time.Now()
			_, err := upstream.client.LookupIP(ctx, "ip", "www.google.com")
			latency := time.Since(start)

			// 更新延迟
			upstream.latency = latency

			// 更新健康状态
			wasHealthy := upstream.healthy
			upstream.healthy = err == nil

			if wasHealthy && !upstream.healthy {
				logger.Warn("上游 DNS 服务器变为不健康",
					"address", upstream.Address,
					"error", err,
					"latency", latency)
			} else if !wasHealthy && upstream.healthy {
				logger.Info("上游 DNS 服务器恢复健康",
					"address", upstream.Address,
					"latency", latency)
			}
		}(us)
	}
}

// StartHealthCheck 启动定期健康检查
func (r *Resolver) StartHealthCheck(interval time.Duration) {
	ticker := time.NewTicker(interval)
	go func() {
		for range ticker.C {
			r.CheckUpstreamHealth()
		}
	}()
	logger.Info("启动上游 DNS 服务器健康检查", "interval", interval)
}

// ClearCache 清除 DNS 缓存
func (r *Resolver) ClearCache() {
	if r.cache != nil {
		r.cache.Clear()
		logger.Info("DNS 缓存已清除")
	}
}

// GetCacheStats 获取缓存统计信息
func (r *Resolver) GetCacheStats() (int, int) {
	if r.cache != nil {
		return r.cache.Size(), r.cache.Capacity()
	}
	return 0, 0
}
